import httpx
import asyncio
from functools import partial
from collections import defaultdict
from typing import Any, List, Tuple, Dict, Callable


def fmt_td(data: dict) -> list | None:
    """格式化数据 tdengine """
    if data.get('code') == 0:
        results = data.get('data')
        if isinstance(results, list) and len(results) == 1:
            return results[0]
        else:
            return results


async def fetch(
    data: Dict | str,
    client: httpx.AsyncClient,
    url: str,
    sem: asyncio.Semaphore,
    func: Callable[..., Any] | None = None,
    is_json: bool = True,
):
    """
    获取数据
    :param data: sql 数据
    :param client: tdengine 客户端
    :param url: tdengine 的 url 地址
    :param sem: 信号量，限制协程数量
    :param func: 函数：用于解析数据
    :param is_json: 是否传递 json 数据到请求服务端
    :return:
    """
    async with sem:
        if is_json:
            r = await client.post(url, json=data)
        else:
            r = await client.post(url, data=data)
        if func is not None:
            return func(r.json())
        return r.json()


async def run(
    data: [list, tuple, str],
    url: str,
    auth: Tuple | None,
    sem_num: int = 20,
    func: Callable[..., Any] | None = None,
    is_json: bool = True,
    mapping: bool = False
) -> Any | Dict:
    """
    协程主函数
    :param data: sql 数据
    :param url: tdengine 的 url 地址
    :param auth: ('username', 'password')
    :param sem_num: 信号量，限制协程数量
    :param func: 函数：用于解析数据
    :param is_json: 是否传递 json 数据到请求服务端
    :param mapping: data 中的数据格式是否字典
    :return:
    """
    sem = asyncio.Semaphore(sem_num)
    async with httpx.AsyncClient(auth=auth) as client:
        if isinstance(data, str):
            return await fetch(data, client, url, sem, func=func, is_json=is_json)
        else:
            if mapping:
                tasks = [fetch(line.sql, client, url, sem, func, is_json)for line in data]
            else:
                tasks = [fetch(line, client, url, sem, func, is_json) for line in data]
        results = await asyncio.gather(*tasks)
        # sql 格式为字典时：处理映射数据并返回
        if mapping:
            return {line.m: results[num] for num, line in enumerate(data)}
        return results


def main(
    data: [list, tuple, str],
    url: str,
    auth: tuple,
    sem_num: int = 50,
    func: Callable[..., Any] | None = None,
    is_json: bool = True,
):
    asyncio.run(run(data, url, auth, sem_num, func, is_json))


async_data = partial(run, is_json=False)


if __name__ == '__main__':
    url = "http://127.0.0.1:6041/rest/sql/nan_steel_shaojie?tz=Asia/Shanghai"
    sqls = ["insert into nan_steel_shaojie.ts_pd_h2_203_temp_mf_h using nan_steel_shaojie.temp_mf_h (eq_uid, part_uid) tags ('pd_h2_203', 'pd_h2_203') values (1743561550000, now, 24.4);",
            "select * from nan_steel_shaojie.ts_pd_h2_203_temp_mf_h tphtmh where tphtmh.ts > '2025-04-02 00:00:00';"]
    # sql = "select (max(kwh) - min(kwh)) v from nan_steel_shaojie.M_UID_10 where ts > '2024-11-01 00:00:00' and ts < '2024-11-25 00:00:00';"
    print(main(sqls, url, ('root', 'taosdata'), func=fmt_td, is_json=False))
